import pandas as pd
import numpy as np

data1 = pd.read_csv(r'data\TRD_Dalyr.csv')
data2 = pd.read_csv(r'data\TRD_Dalyr1.csv')
print(data1.shape)
print(data2.shape)
data = pd.concat([data1, data2])  #
print(data)
data['LogDretwd'] = np.log(data['Dretwd'] + 1)
print(data)

stock_codes = data['Stkcd'].unique()
print(stock_codes)

df = pd.DataFrame(index=data['Trddt'].unique())
for stkcd in stock_codes:
    temp = data[data['Stkcd'] == stkcd]
    if temp.shape[0] == df.shape[0]:
        df[stkcd] = temp['LogDretwd'].values
print(df)

df = df.iloc[:, :500]  # 选取前500只股票，减少运行时间

mean_returns = df.mean()
print(mean_returns)
std_returns = df.std()
print(std_returns)
cov_matrix = df.cov()
print(cov_matrix)


def getRisk(x, cov_matrix):
    return np.sqrt(np.dot(x, np.dot(cov_matrix, x)))


def getReturn(x, mean_returns):
    return np.dot(x, mean_returns)


stock_num = df.shape[1]
print(stock_num)

initiate_x = [1 / stock_num] * stock_num

from scipy import optimize

frontier = list()
for target_return in np.linspace(mean_returns.min(), mean_returns.max(), 20):
    print(target_return)
    bounds = tuple((0, 1) for _ in range(stock_num))
    constraints = ({'type': 'eq', 'fun': lambda x: target_return - getReturn(x, mean_returns)},
                   {'type': 'eq', 'fun': lambda x: 1 - sum(x)}
                   )
    minimize_risk = optimize.minimize(getRisk,
                                      initiate_x,
                                      args=cov_matrix,
                                      method='SLSQP',
                                      bounds=bounds,
                                      constraints=constraints)
    initiate_x = minimize_risk.x
    frontier.append((target_return, minimize_risk.fun, minimize_risk.x.tolist()))
print(frontier)

import matplotlib.pyplot as plt

fig, ax = plt.subplots()
ax.plot(std_returns, mean_returns, '.')
ax.plot([x[1] for x in frontier], [x[0] for x in frontier])
plt.xlabel('risk')
plt.ylabel('return')
plt.show()
